package name.zhengkeli.cloud.auth

import com.auth0.jwt.JWT
import com.auth0.jwt.algorithms.Algorithm
import com.auth0.jwt.exceptions.JWTVerificationException
import io.ktor.http.*
import io.ktor.server.application.*
import io.ktor.server.auth.*
import io.ktor.server.config.*
import io.ktor.server.request.*
import io.ktor.server.response.*
import io.ktor.server.routing.*
import kotlinx.coroutines.runBlocking
import kotlinx.datetime.Clock
import kotlinx.datetime.Instant
import kotlinx.datetime.toJavaInstant
import kotlinx.datetime.toKotlinInstant
import kotlinx.serialization.Serializable
import name.zhengkeli.cloud.auth.AuthSchema.PasswordsTable
import name.zhengkeli.cloud.auth.AuthSchema.UsernamesTable
import name.zhengkeli.cloud.auth.AuthSchema.UsersTable
import org.jetbrains.exposed.sql.*
import java.security.MessageDigest
import java.util.*
import kotlin.time.Duration


@Serializable
data class RegisterUserResponse(val userId: String, val token: String)

class RegisterUserNotAllowedException() : ServerException(HttpStatusCode.Forbidden)


@Serializable
data class AddUsernameRequest(val username: String)

@Serializable
data class AddUsernameResponse(val usernameId: String)

class UsernameIdNotFoundException() : ServerException(HttpStatusCode.Unauthorized)


@Serializable
data class AddPasswordRequest(val password: String)

@Serializable
data class AddPasswordResponse(val passwordId: String)

class PasswordIdNotFoundException() : ServerException(HttpStatusCode.Unauthorized)


@Serializable
data class LoginRequest(val username: String, val password: String)

@Serializable
data class LoginResponse(val userId: String, val token: String)

class UsernameNotFoundException() : ServerException(HttpStatusCode.Unauthorized)

class UserIdNotFoundException() : ServerException(HttpStatusCode.Unauthorized)

class PasswordNotSetException() : ServerException(HttpStatusCode.Unauthorized)

class PasswordNotMatchException() : ServerException(HttpStatusCode.Unauthorized)


data class Principal(val userId: UUID)

class UnauthorizedException() : ServerException(HttpStatusCode.Unauthorized)

inline fun <reified P : Any> ApplicationCall.principalOrThrow(): P {
    return principal<P>() ?: throw UnauthorizedException()
}


class PathParameterNotSetException() : ServerException(HttpStatusCode.BadRequest)

class PathParameterIllegalException() : ServerException(HttpStatusCode.BadRequest)

fun ApplicationCall.pathParameterOrThrow(name: String): String {
    return parameters[name] ?: throw PathParameterNotSetException()
}

inline fun <T> ApplicationCall.pathParameterOrThrow(name: String, block: (String) -> T): T {
    val value = pathParameterOrThrow(name)
    val result = runCatching { block(value) }
    return result.getOrElse { throw PathParameterIllegalException() }
}


fun Application.configureAuthService(config: ApplicationConfig) {
    val databaseUrl = config.property("database.url").getString()
    val databaseDriver = config.property("database.driver").getString()
    val databaseUser = config.property("database.user").getString()
    val databasePassword = config.property("database.password").getString()

    val database = runBlocking {
        Database.connect(
            url = databaseUrl,
            driver = databaseDriver,
            user = databaseUser,
            password = databasePassword,
        ).initialize()
    }

    val random = kotlin.random.Random.Default
    val serverSecret = config.property("serverSecret").getString().toByteArray(Charsets.UTF_8)
    val tokenDuration = config.property("tokenDuration").getString().let { Duration.parse(it) }
    val allowRegister = config.property("allowRegister").getString().toBoolean()

    authentication {
        bearer {
            authenticate { (token) ->
                val jwt = JWT.decode(token)

                val nowInstant = Clock.System.now()
                if (nowInstant > jwt.expiresAtAsInstant.toKotlinInstant())
                    return@authenticate null // token expired
                if (nowInstant < jwt.notBeforeAsInstant.toKotlinInstant())
                    return@authenticate null // token not valid yet

                val userId = UUID.fromString(jwt.subject)
                val userSecret = database.transaction {
                    UsersTable.select(UsersTable.secret)
                        .where { UsersTable.id eq userId }
                        .andWhere { UsersTable.deleteTime.isNull() }
                        .map { it[UsersTable.secret] }
                        .singleOrNull()
                } ?: return@authenticate null // user not found

                try {
                    val secret = userSecret + serverSecret
                    JWT.require(Algorithm.HMAC256(secret)).build().verify(jwt)
                } catch (_: JWTVerificationException) {
                    return@authenticate null // signature invalid
                }

                return@authenticate Principal(userId)
            }
        }
    }

    routing {
        post("/users") {
            if (!allowRegister)
                throw RegisterUserNotAllowedException()

            val userSecret = random.nextBytes(UsersTable.SECRET_SIZE)

            val nowInstant = Clock.System.now()
            val userId = database.transaction {
                UsersTable.insertAndGetId {
                    it[UsersTable.secret] = userSecret
                    it[UsersTable.createTime] = nowInstant
                    it[UsersTable.updateTime] = nowInstant
                }.value
            }

            val token = issueToken(userId, userSecret, serverSecret, tokenDuration, nowInstant)

            call.respond(RegisterUserResponse(userId.toString(), token))
        }
        post("/login") {
            val (username, password) = call.receive<LoginRequest>()

            val userId = database.transaction {
                UsernamesTable.select(UsernamesTable.userId)
                    .where { UsernamesTable.username eq username }
                    .andWhere { UsernamesTable.deleteTime.isNull() }
                    .map { it[UsernamesTable.userId] }
                    .singleOrNull()
            } ?: throw UsernameNotFoundException()

            val userSecret = database.transaction {
                UsersTable.select(UsersTable.secret)
                    .where { UsersTable.id eq userId }
                    .andWhere { UsersTable.deleteTime.isNull() }
                    .map { it[UsersTable.secret] }
                    .singleOrNull()
            } ?: throw UserIdNotFoundException()

            val userPasswords = database.transaction {
                PasswordsTable.select(PasswordsTable.salt, PasswordsTable.hash)
                    .where { PasswordsTable.userId eq userId }
                    .andWhere { PasswordsTable.deleteTime.isNull() }
                    .map { Pair(it[PasswordsTable.salt], it[PasswordsTable.hash]) }
            }
            if (userPasswords.isEmpty())
                throw PasswordNotSetException()
            if (!userPasswords.any { (salt, hash) -> verifyPassword(password, salt, hash) })
                throw PasswordNotMatchException()

            val nowInstant = Clock.System.now()
            val token = issueToken(userId, userSecret, serverSecret, tokenDuration, nowInstant)
            call.respond(LoginResponse(userId.toString(), token))
        }
        authenticate {
            post("/verify") {
                call.respond(HttpStatusCode.OK)
            }
            post("/usernames") {
                val (userId) = call.principalOrThrow<Principal>()
                val (username) = call.receive<AddUsernameRequest>()

                val nowInstant = Clock.System.now()
                val usernameId = database.transaction {
                    UsernamesTable.insertAndGetId {
                        it[UsernamesTable.userId] = userId
                        it[UsernamesTable.username] = username
                        it[UsernamesTable.createTime] = nowInstant
                    }.value
                }

                call.respond(AddUsernameResponse(usernameId.toString()))
            }
            delete("/usernames/{usernameId}") {
                val (userId) = call.principalOrThrow<Principal>()
                val usernameId = call.pathParameterOrThrow("usernameId", { UUID.fromString(it) })

                val nowInstant = Clock.System.now()
                database.transaction {
                    val ownerUserId = UsernamesTable
                        .select(UsernamesTable.userId)
                        .where { UsernamesTable.id eq usernameId }
                        .andWhere { UsernamesTable.deleteTime.isNull() }
                        .map { it[UsernamesTable.userId] }
                        .singleOrNull()
                        ?: throw UsernameIdNotFoundException()

                    if (ownerUserId != userId)
                        throw UnauthorizedException()

                    UsernamesTable.update({
                        UsernamesTable.id eq usernameId
                    }) { it[UsernamesTable.deleteTime] = nowInstant }
                }

                call.respond(HttpStatusCode.OK)
            }
            post("/passwords") {
                val (userId) = call.principalOrThrow<Principal>()
                val (password) = call.receive<AddPasswordRequest>()

                val salt = random.nextBytes(PasswordsTable.SALT_SIZE)
                val hash = encryptPassword(password, salt)

                val nowInstant = Clock.System.now()
                val passwordId = database.transaction {
                    PasswordsTable.insertAndGetId {
                        it[PasswordsTable.userId] = userId
                        it[PasswordsTable.salt] = salt
                        it[PasswordsTable.hash] = hash
                        it[PasswordsTable.createTime] = nowInstant
                    }
                }

                call.respond(AddPasswordResponse(passwordId.toString()))
            }
            delete("/passwords/{passwordId}") {
                val (userId) = call.principalOrThrow<Principal>()
                val passwordId = call.pathParameterOrThrow("passwordId", { UUID.fromString(it) })

                val nowInstant = Clock.System.now()
                database.transaction {
                    val ownerUserId = PasswordsTable
                        .select(PasswordsTable.userId)
                        .where { PasswordsTable.id eq passwordId }
                        .andWhere { PasswordsTable.deleteTime.isNull() }
                        .map { it[PasswordsTable.userId] }
                        .singleOrNull()
                        ?: throw PasswordIdNotFoundException()

                    if (ownerUserId != userId)
                        throw UnauthorizedException()

                    PasswordsTable.update({
                        PasswordsTable.id eq passwordId
                    }) { it[PasswordsTable.deleteTime] = nowInstant }
                }

                call.respond(HttpStatusCode.OK)
            }
        }
    }
}

fun encryptPassword(password: String, salt: ByteArray): ByteArray {
    val md = MessageDigest.getInstance("SHA-256")
    md.update(password.toByteArray(Charsets.UTF_8))
    md.update(salt)
    return md.digest()
}

fun verifyPassword(password: String, salt: ByteArray, hash: ByteArray): Boolean {
    val md = MessageDigest.getInstance("SHA-256")
    md.update(password.toByteArray(Charsets.UTF_8))
    md.update(salt)
    return md.digest().contentEquals(hash)
}

fun issueToken(
    userId: UUID,
    userSecret: ByteArray,
    serverSecret: ByteArray,
    tokenDuration: Duration,
    nowInstant: Instant
): String {
    val expInstant = nowInstant + tokenDuration
    val secret = userSecret + serverSecret
    return JWT.create()
        .withSubject(userId.toString())
        .withNotBefore(nowInstant.toJavaInstant())
        .withExpiresAt(expInstant.toJavaInstant())
        .withClaim("use", "login")
        .sign(Algorithm.HMAC256(secret))
}